use std::cmp::Ordering;

#[derive(Debug)]
enum HeapType {
    Min,
    Max,
}

struct BinaryHeap<T> {
    data: Vec<T>,
    heap_type: HeapType,
}

impl<T: Ord + Clone> BinaryHeap<T> {
    fn new(heap_type: HeapType) -> Self {
        BinaryHeap {
            data: Vec::new(),
            heap_type,
        }
    }

    fn insert(&mut self, value: T) {
        self.data.push(value);
        let mut index = self.data.len() - 1;

        while index > 0 {
            let parent_index = (index - 1) / 2;
            match self.heap_type {
                HeapType::Min => {
                    if self.data[index].cmp(&self.data[parent_index]) == Ordering::Less {
                        self.data.swap(index, parent_index);
                        index = parent_index;
                    } else {
                        break;
                    }
                }
                HeapType::Max => {
                    if self.data[index].cmp(&self.data[parent_index]) == Ordering::Greater {
                        self.data.swap(index, parent_index);
                        index = parent_index;
                    } else {
                        break;
                    }
                }
            }
        }
    }

    fn pop(&mut self) -> Option<T> {
        if self.data.is_empty() {
            return None;
        }

        let result = self.data[0].clone();
        let last_value = self.data.pop().unwrap();

        if !self.data.is_empty() {
            self.data[0] = last_value;
            let mut index = 0;

            loop {
                let left_child_index = index * 2 + 1;
                let right_child_index = index * 2 + 2;
                let mut next_index = index;

                if left_child_index < self.data.len() {
                    match self.heap_type {
                        HeapType::Min => {
                            if self.data[left_child_index].cmp(&self.data[next_index]) == Ordering::Less {
                                next_index = left_child_index;
                            }
                        }
                        HeapType::Max => {
                            if self.data[left_child_index].cmp(&self.data[next_index]) == Ordering::Greater {
                                next_index = left_child_index;
                            }
                        }
                    }
                }

                if right_child_index < self.data.len() {
                    match self.heap_type {
                        HeapType::Min => {
                            if self.data[right_child_index].cmp(&self.data[next_index]) == Ordering::Less {
                                next_index = right_child_index;
                            }
                        }
                        HeapType::Max => {
                            if self.data[right_child_index].cmp(&self.data[next_index]) == Ordering::Greater {
                                next_index = right_child_index;
                            }
                        }
                    }
                }

                if index == next_index {
                    break;
                } else {
                    self.data.swap(index, next_index);
                    index = next_index;
                }
            }
        }

        Some(result)
    }

    fn num_entries(&self) -> usize {
        self.data.len()
    }
}

fn main() {
    let mut min_heap = BinaryHeap::new(HeapType::Min);
    min_heap.insert(5);
    min_heap.insert(3);
    min_heap.insert(8);
    println!("{:?}", min_heap.pop()); // Some(3)
    println!("{:?}", min_heap.num_entries()); // 2

    let mut max_heap = BinaryHeap::new(HeapType::Max);
    max_heap.insert(5);
    max_heap.insert(3);
    max_heap.insert(8);
    println!("{:?}", max_heap.pop()); // Some(8)
    println!("{:?}", max_heap.num_entries()); // 2
}
